package com.ktjiaoyu.travelaiagent.mcp;

import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.model.function.FunctionCallback;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;

/**
 * MCP服务工具代理类，解决ToolContext异常
 */
public class McpToolCallbackProxy implements InvocationHandler {
    // 代理目标对象
    private final FunctionCallback target;

    public McpToolCallbackProxy(FunctionCallback target) {
        this.target = target;
    }
    /** 拦截带Tool Context参数的call方法进行特殊处理 */
    @Override
    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
        // 这个默认认为类名包含MCP，就是MCP服务注册的工具
        if (target.getClass().getSimpleName().toLowerCase().contains("mcp")
                && method.getName().equals("call") // 调用了call方法
                // 并且第二个参数是ToolContext类型
                && args.length == 2 && args[1].getClass().equals(ToolContext.class)) {
            // 满足上诉条件，直接调用call(String)，避免抛异常
            return target.call(args[0].toString());
        }
        return method.invoke(target, args);
    }

    /** 用于对工具产生代理对象 */
    public static FunctionCallback[] proxyAll(FunctionCallback... callbacks) {
        // 创建代理对象数组
        FunctionCallback[] proxyArray = new FunctionCallback[callbacks.length];
        for (int i = 0; i < callbacks.length; i++) {
            // 需代理的对象
            FunctionCallback callback = callbacks[i];
            // 创建代理对象
            proxyArray[i] = (FunctionCallback) Proxy.newProxyInstance(
                    callback.getClass().getClassLoader(),
                    callback.getClass().getInterfaces(),
                    new McpToolCallbackProxy(callback));
        }
        return proxyArray;
    }
}
